from keras.callbacks import ModelCheckpoint
import warnings
import mlflow
import mlflow.keras
from image_pyfunc import log_model, KerasImageClassifierPyfunc


class BoltNutModelCheckpoint(ModelCheckpoint):
    """ 
    Keras callback for logging metrics and final model with MLflow.

    Metrics are logged after every epoch. The logger keeps track of the best model based on the
    validation metric. At the end of the training, the best model is logged with MLflow.
    """

    def __init__(self, model_to_save, json_path, **kwargs):
        super(BoltNutModelCheckpoint, self).__init__(**kwargs)
        self.model_to_save = model_to_save
        self.json_path = json_path

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        self.epochs_since_last_save += 1
        if self.epochs_since_last_save >= self.period:
            self.epochs_since_last_save = 0
            filepath = self.filepath.format(epoch=epoch + 1, **logs)
            if self.save_best_only:
                current = logs.get(self.monitor)
                if current is None:
                    warnings.warn('Can save best model only with %s available, '
                                  'skipping.' % (self.monitor), RuntimeWarning)
                else:
                    if self.monitor_op(current, self.best):
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                                  ' saving model to %s'
                                  % (epoch + 1, self.monitor, self.best,
                                     current, filepath))
                        self.best = current
                        if self.save_weights_only:
                            self.model_to_save.save_weights(filepath, overwrite=True)
                        else:
                            print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
                            mlflow.log_metric(key="loss", value=float(current), step=epoch + 1)
                            log_model(keras_model=self.model_to_save, json_path=self.json_path)
                    else:
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s did not improve from %0.5f' %
                                  (epoch + 1, self.monitor, self.best))
            else:
                if self.verbose > 0:
                    print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
                if self.save_weights_only:
                    self.model_to_save.save_weights(filepath, overwrite=True)
                else:
                    self.model_to_save.save(filepath, overwrite=True)

        super(BoltNutModelCheckpoint, self).on_batch_end(epoch, logs)

    def on_train_end(self, logs=None):
        """
        Log the best model with MLflow and evaluate it on the train and validation data so that the
        metrics stored with MLflow reflect the logged model.
        """
        # log_model(keras_model=self.model_to_save, **self._pyfunc_params)